# 开发时间: 2024/2/18 1:11

import math

pts = []  # 点集，任意维度的点集
targetPt = []# 目标点，任意维度的点


class Node():
    def __init__(self, pt, leftBranch, rightBranch, dimension):
        self.pt = pt
        self.leftBranch = leftBranch
        self.rightBranch = rightBranch
        self.dimension = dimension


class KDTree():
    def __init__(self, data):
        self.nearestPt = None
        self.nearestDis = math.inf

    def createKDTree(self, currPts, dimension):
        if (len(currPts) == 0):
            return None
        mid = self.calMedium(currPts)
        sortedData = sorted(currPts, key=lambda x: x[dimension])
        leftBranch = self.createKDTree(sortedData[:mid], self.calDimension(dimension))
        rightBranch = self.createKDTree(sortedData[mid + 1:], self.calDimension(dimension))
        return Node(sortedData[mid], leftBranch, rightBranch, dimension)

    def calMedium(self, currPts):
        return len(currPts) // 2

    def calDimension(self, dimension):  # 区别就在于这里，几维就取余几
        return (dimension + 1) % len(targetPt)

    def calDistance(self, p0, p1):
        return math.sqrt((p0[0] - p1[0]) ** 2 + (p0[1] - p1[1]) ** 2)

    def getNearestPt(self, root, targetPt):
        self.search(root, targetPt)
        return self.nearestPt, self.nearestDis

    def search(self, node, targetPt):
        if node == None:
            return
        dist = node.pt[node.dimension] - targetPt[node.dimension]
        if (dist > 0):  # 目标点在节点的左侧或上侧
            self.search(node.leftBranch, targetPt)
        else:
            self.search(node.rightBranch, targetPt)
        tempDis = self.calDistance(node.pt, targetPt)
        if (tempDis < self.nearestDis):
            self.nearestDis = tempDis
            self.nearestPt = node.pt
        # 回溯
        if (self.nearestDis > abs(dist)):
            if (dist > 0):
                self.search(node.rightBranch, targetPt)
            else:
                self.search(node.leftBranch, targetPt)


if __name__ == "__main__":
    kdtree = KDTree(pts)
    root = kdtree.createKDTree(pts, 0)
    pt, minDis = kdtree.getNearestPt(root, targetPt)
    print("最近的点是", pt, "最小距离是", str(minDis))